from __future__ import division
import itertools
from vtk import *
from pylab import *

class vtkSpectrogram:
    def __init__(self):
        self.renderer = vtkRenderer()
        self.renderer.SetBackground(1.0, 1.0, 1.0)
        self.renderWindow = vtkRenderWindow()
        self.renderWindow.AddRenderer(self.renderer)
        self.interactorStyle = vtkInteractorStyleTrackballCamera()
        self.renderWindowInteractor = vtkRenderWindowInteractor()
        self.renderWindowInteractor.SetRenderWindow(self.renderWindow)
        self.renderWindowInteractor.SetInteractorStyle(self.interactorStyle)
        
#        self.renderer.GetActiveCamera().Azimuth(30)
#        self.renderer.GetActiveCamera().Elevation(30)
        self.renderer.GetActiveCamera().ParallelProjectionOn() 
        self.renderer.ResetCamera()

    def start(self):
        self.renderWindowInteractor.Start()
    def addLine(self, x, y, z, radius, color):
        points = vtkPoints()
        map(lambda x_i, y_i, z_i: points.InsertNextPoint(x_i, y_i, z_i), x, y, z)

        lines = vtkCellArray()
        lines.InsertNextCell(points.GetNumberOfPoints())
        map(lambda i: lines.InsertCellPoint(i), xrange(points.GetNumberOfPoints()))
        
        polyData = vtkPolyData()
        polyData.SetPoints(points)
        polyData.SetLines(lines)
        
        tubeRadius = vtkDoubleArray()
        tubeRadius.SetName('TubeRadius')
        tubeRadius.SetNumberOfValues(len(radius))
        map(lambda (i, radius_i): tubeRadius.SetValue(i, radius_i), enumerate(radius))
        polyData.GetPointData().AddArray(tubeRadius)
        polyData.GetPointData().SetActiveScalars('TubeRadius')
        
        colors = vtkUnsignedCharArray()
        colors.SetName('Colors')
        colors.SetNumberOfTuples(len(color))
        colors.SetNumberOfComponents(1)
        map(lambda (i, color_i): colors.SetValue(i, int(color_i)), enumerate(color))
        polyData.GetPointData().AddArray(colors)
        
        tube = vtkTubeFilter()
        tube.SetInput(polyData)
        tube.SetNumberOfSides(30)
#        tube.SetDefaultNormal(0, 0, 1)
        tube.SetVaryRadiusToVaryRadiusByAbsoluteScalar()
        tube.CappingOn()
        
#        ribbon = vtkRibbonFilter()
#        ribbon.SetInput(polyData)
#        ribbon.VaryWidthOn()
#        ribbon.SetWidthFactor(10)
#        ribbon.SetWidth(0.3)
        
        mapper = vtkPolyDataMapper()
        mapper.SetInputConnection(tube.GetOutputPort())
#        mapper.SetInputConnection(ribbon.GetOutputPort())
#        mapper.ScalarVisibilityOn()
        mapper.SetScalarModeToUsePointFieldData()
        mapper.SelectColorArray('Colors')
        mapper.Update()
        mapper.StaticOn()
#        mapper.ImmediateModeRenderingOn()
        
        actor = vtkActor()
        actor.SetMapper(mapper)
#        actor.GetProperty().SetAmbient(1)
#        actor.GetProperty().SetDiffuse(0)
#        actor.GetProperty().SetSpecular(0)
        
        self.renderer.AddActor(actor)
        return actor
    
if __name__ == '__main__':
    test = vtkSpectrogram()
    t = linspace(0,1,500)
    test.addLine(t, t**2, zeros(500), t/10, t*255)
    test.addLine(-t, t**2, zeros(500), t/10, t*255)
    test.start()

